import { assetManager, BufferAsset, js } from 'cc';
import wonnx from './wonnx.js';
import { EDITOR } from 'cc/env';

export class Input {
	private _input: wonnx.Input = null!;
	public get input(): wonnx.Input {
		return this._input;
	}
	public set input(value: wonnx.Input) {
		this._input = value;
	}
	constructor() {
		this._input = new wonnx.Input();
	}

	insert(input_name: string, value: Float32Array): void {
		this._input.insert(input_name, value);
	}

	free(): void {
		this._input.free();
	}
}

export class Session {
	private _session: wonnx.Session = null!;
	free(): void {
		this._session.free();
	}
	/**
	 * @param {Uint8Array} bytes
	 * @returns {Promise<Session>}
	 */
	static fromBytes(bytes: Uint8Array): Promise<Session> {
		return new Promise<Session>((resolve, reject) => {
			wonnx.Session.fromBytes(bytes)
				.then((session) => {
					const s = new Session();
					s._session = session;
					resolve(s);
				})
				.catch((error) => {
					reject(error);
				});
		});
	}
	/**
	 * @param {Input} input
	 * @returns {Promise<any>}
	 */
	run(input: Input): Promise<any> {
		return new Promise<any>((resolve, reject) => {
			this._session
				.run(input.input)
				.then((result) => {
					resolve(result);
				})
				.catch((error) => {
					reject(error);
				});
		});
	}
}

const moduleFactory = wonnx.moduleFactory;

async function loadWasmBinary(wasmBinaryUuid: string) {
	let wasmBinary: Uint8Array | undefined;
	if (!EDITOR) {
		const bufferAsset = await new Promise<BufferAsset | null>((resolve, reject) => {
			assetManager.loadBundle('mlagentResource', (error, bundle) => {
				if (error) {
					console.error(error);
				} else {
					bundle.load<BufferAsset>('WebAssembly/wonnx_bg', (error, asset) => {
						if (error) {
							console.error(error);
							reject(null);
						} else {
							resolve(asset);
						}
					});
				}
			});
		});
		wasmBinary = bufferAsset == null ? undefined : new Uint8Array(bufferAsset.buffer());
	} else {
		const wasmFilePath = await Editor.Message.request('asset-db', 'query-path', wasmBinaryUuid);
		console.info(`wasmFilePath: ${wasmFilePath}`);
		if (!wasmFilePath) {
			throw new Error(`'Failed to load WebAssembly Binary ${wasmBinaryUuid}`);
		}
		const buffer = require('fs').readFileSync(wasmFilePath);
		wasmBinary = new Uint8Array(buffer.buffer, buffer.byteOffset, buffer.byteLength);
	}
	return wasmBinary;
}

async function loadWONNXModule() {
	console.log(`Loading WONNX  Core.`);

	const myModulePromiseCacheKey = Symbol.for('WONNXModulePromise');
	const global = globalThis as unknown as {
		[myModulePromiseCacheKey]: ReturnType<typeof moduleFactory>;
	};
	const modulePromise = (global[myModulePromiseCacheKey] ??= (async () => {
		const wasmBinaryUuid = 'ddac461b-4961-48b0-8974-b088cee3bdf2';
		const wasmBinary = await loadWasmBinary(wasmBinaryUuid);
		const index = await wonnx.moduleFactory(wasmBinary!);
		if (!index) {
			throw new Error(`'Failed to get WONNX module factory.`);
		}
		return index;
	})());

	return await modulePromise;
}

await loadWONNXModule();
